# -*-coding:utf-8 -*-
# author: leemon
# datetime: 2020/8/24 13:41

import requests
import json


postHeader = {'Content-Type': 'application/json'}
pritunl_url = {
    'authSessionAPI': '/auth/session',
    'userAPI': '/user',
    'downloadProfileLinkAPI': '/key',
    'getCsrfTokenAPI': '/state',
}
return_tmp = {'status': True, 'msg': ''}


class Pritunl:
    def __init__(self, url, serverid, admin, password):
        self.url = url
        self.serverid = serverid
        self.admin = admin
        self.password = password
        requests.packages.urllib3.disable_warnings()
        self.req = requests.session()
        self.header = {}
        self.get_token()

    def get_token(self):
        post_login_data = {
            "username": self.admin,
            "password": self.password
        }
        try:
            # do login
            ret = self.req.post('{}{}'.format(self.url, pritunl_url['authSessionAPI']), data=json.dumps(post_login_data), headers=postHeader, verify=False)
            # do request token
            res = self.req.get('{}{}'.format(self.url, pritunl_url['getCsrfTokenAPI']), headers=postHeader, verify=False)
        except:
            return False
        else:
            res_json = json.loads(res.text)

            token = res_json.get('csrf_token', False)
            if token:
                self.header = {'Csrf-Token': token}
            return token

    def useradd(self, username, password='123456'):
        create_user_url = '{}{}/{}'.format(self.url, pritunl_url['userAPI'], self.serverid)
        usersearch_res = self.usersearch(username)
        if usersearch_res.get('users', False):
            return_tmp['status'] = False
            return_tmp['msg'] = 'User already exist: {}'.format(username)
            return return_tmp
        user_body = {
            'name': username,
            'email': '',
            'auth_type': 'local',
            'bypass_secondary': None,
            'client_to_client': None,
            'network_links': None,
            'dns_servers': None,
            'dns_suffix': None,
            'port_forwarding': None,
            'pin': None
        }

        postHeader.update(self.header)
        ret = self.req.post(
            create_user_url,
            data=json.dumps(user_body),
            headers=postHeader,
            verify=False
        )
        if ret.status_code != 200:
            return_tmp['status'] = False
            return_tmp['msg'] = 'User [{}] create failed: {}'.format(username, ret.content)
            return return_tmp
        res_json = json.loads(ret.content)
        userid = res_json[0].get('id')
        profile = self.get_profile_link(userid)
        if not profile:
            return_tmp['status'] = False
            return_tmp['msg'] = 'User [{}] profile get failed!'.format(username)
            return return_tmp
        profile_url = '{}{}'.format(self.url, profile.get('view_url'))
        result = {'status': True, 'username': username, 'password': None, 'view_url': profile_url, 'msg': '', 'type': 'pritunl'}
        return result

    def userdel(self, username):
        usersearch_res = self.usersearch(username)
        users = usersearch_res.get('users', False)
        if len(users) == 0:
            status = False
            msg = 'User not found: {}'.format(username)
        elif len(users) == 1:
            # do delete
            userid = users[0].get('id')
            del_user_url = '{}{}/{}/{}'.format(self.url, pritunl_url['userAPI'], self.serverid, userid)
            res = self.req.delete(
                del_user_url,
                headers=self.header,
                verify=False
            )
            if res.status_code != 200:
                status = False
                msg = 'Call pritunl api error when delete user: {}'.format(username)
            else:
                status = True
                msg = 'User delete successfully: {}'.format(username)
        else:
            status = False
            msg = 'Found {} users by the username {}, Please use the admin page to delete user.'.format(len(users), username)

        return_tmp['status'] = status
        return_tmp['msg'] = msg
        return return_tmp

    def userdelbyid(self, userid):
            # do delete
        del_user_url = '{}{}/{}/{}'.format(self.url, pritunl_url['userAPI'], self.serverid, userid)
        res = self.req.delete(
            del_user_url,
            headers=self.header,
            verify=False
        )
        if res.status_code != 200:
            status = False
            msg = 'Call pritunl api error when delete user: {}'.format(userid)
        else:
            status = True
            msg = 'User delete successfully: {}'.format(userid)

        return_tmp['status'] = status
        return_tmp['msg'] = msg
        return return_tmp

    def usersearch(self, username):
        res = self.req.get(
            '{}{}/{}?search={}'.format(self.url, pritunl_url['userAPI'], self.serverid, username),
            headers=self.header,
            verify=False
            )

        return json.loads(res.content)

    def get_all_users(self):
        page = 0
        users = []
        users_list_tmp = []
        while True:
            res = self.req.get(
                '{}{}/{}?page={}'.format(self.url, pritunl_url['userAPI'], self.serverid, page),
                headers=self.header,
                verify=False
                )
            if res.status_code != 200:
                return False
            res_json = json.loads(res.content)
            page_total = res_json.get('page_total', False)
            users_list_tmp.extend(res_json.get('users'))
            if page_total and page < page_total:
                page += 1
            else:
                break

        for user in users_list_tmp:
            user_dict = {'userid': '', 'name': '', 'email': '', 'organization': '', 'disabled': False}
            if user['type'] == 'client':
                user_dict['userid'] =  user['id']
                user_dict['name'] = user['name']
                user_dict['email'] = user['email']
                user_dict['organization'] = user['organization']
                user_dict['disabled'] = user['disabled']
                users.append(user_dict)
        return users

    def get_profile_link(self, userid):
        get_profile_link_url = '{}{}/{}/{}'.format(self.url, pritunl_url['downloadProfileLinkAPI'], self.serverid, userid)
        res = self.req.get(
            get_profile_link_url,
            headers=self.header,
            verify=False
        )
        if res.status_code != 200:
            return False
        return json.loads(res.content)


if __name__ == '__main__':
    PRITUNL_WEB_URL = 'https://192.168.42.150'
    SERVER_ID = '5f2d136017bc45bf418b8db9'
    PRITUNL_ADMIN_USER = 'pritunl'
    PRITUNL_ADMIN_PASSWORD = 'DdZl1QNgEG2F'
    pritunl = Pritunl(PRITUNL_WEB_URL, SERVER_ID, PRITUNL_ADMIN_USER, PRITUNL_ADMIN_PASSWORD)
    # all_user = pritunl.get_all_users()